#include "common.h"

// 确保完整发送数据 
int send_all(int sockfd, const void *buf, size_t len, struct sockaddr_in *addr, socklen_t addr_len) 
{
    ssize_t sent = sendto(sockfd, buf, len, 0, (struct sockaddr*)addr, addr_len);
    if (sent != (ssize_t)len) 
    {
        return -1; // 发送错误
    }
    
    return 0;
}

// 确保完整接收数据 
int recv_all(int sockfd, void *buf, size_t len, struct sockaddr_in *addr, socklen_t *addr_len) 
{
    ssize_t received = recvfrom(sockfd, buf, len, 0, (struct sockaddr*)addr, addr_len);
    if (received != (ssize_t)len) 
    {
        return -1; // 接收错误
    }
    
    return 0;
}

int send_file(int sockfd, const char *filepath, struct sockaddr_in *addr, socklen_t addr_len) 
{
    FILE *file = NULL;
    int ret = -1;
    
    // 检查文件是否存在且是普通文件
    struct stat st;
    if (stat(filepath, &st) != 0) 
    {
        fprintf(stderr, "stat error for %s: %s\n", filepath, strerror(errno));
        goto cleanup;
    }
    
    if (!S_ISREG(st.st_mode)) 
    {
        fprintf(stderr, "%s is not a regular file\n", filepath);
        goto cleanup;
    }
    
    file = fopen(filepath, "rb");
    if (!file) 
    {
        fprintf(stderr, "fopen error for %s: %s\n", filepath, strerror(errno));
        goto cleanup;
    }

    uint64_t file_size = st.st_size;
    printf("File size: %lu bytes\n", file_size);

    // 提取文件名
    char filename[256];
    strncpy(filename, basename(filepath), sizeof(filename)-1);
    filename[sizeof(filename)-1] = '\0';

    printf("%s sending.....\n", filename);
    
    // 发送文件名长度和文件名
    uint32_t name_len = strlen(filename) + 1; // 包含终止符
    uint32_t net_name_len = htonl(name_len);
    if (send_all(sockfd, &net_name_len, sizeof(net_name_len), addr, addr_len) < 0) 
    {
        perror("send filename length error");
        goto cleanup;
    }
    
    if (send_all(sockfd, filename, name_len, addr, addr_len) < 0) 
    {
        perror("send filename error");
        goto cleanup;
    }
    
    // 发送文件大小（转换为网络字节序）
    uint64_t net_file_size = htobe64(file_size);
    if (send_all(sockfd, &net_file_size, sizeof(net_file_size), addr, addr_len) < 0) 
    {
        perror("send file size error");
        goto cleanup;
    }
    
    // 发送文件内容 (分块发送)
    char sendbuff[1400]; // UDP适合较小的数据包
    size_t readlen;
    uint64_t total_sent = 0;
    
    while ((readlen = fread(sendbuff, 1, sizeof(sendbuff), file)) > 0) 
    {
        if (send_all(sockfd, sendbuff, readlen, addr, addr_len) < 0) 
        {
            perror("send file content error");
            goto cleanup;
        }
        total_sent += readlen;
        
        // 显示进度
        printf("Sent: %lu/%lu bytes (%.2f%%)\r", 
               total_sent, file_size, (double)total_sent/file_size*100);
        fflush(stdout);
        
        // UDP需要短暂延迟以避免丢包
        usleep(1000);
    }
    
    printf("\n%s sent successfully!\n", filename);
    ret = 0; // 成功

cleanup:
    if (file) 
    {
        fclose(file);
    }
    return ret;
}

int recv_file(int sockfd, const char *save_dir, struct sockaddr_in *addr, socklen_t addr_len) 
{
    FILE *file = NULL;
    int ret = -1;
    char filename[256] = {0};	
    
    // 接收文件名长度
    uint32_t net_name_len;
    if (recv_all(sockfd, &net_name_len, sizeof(net_name_len), addr, &addr_len) < 0) 
    {
        perror("recv filename length error");
        goto cleanup;
    }
    
    uint32_t name_len = ntohl(net_name_len);
    if (name_len >= sizeof(filename)) 
    {
        fprintf(stderr, "Filename too long: %u\n", name_len);
        goto cleanup;
    }
    
    // 接收文件名
    if (recv_all(sockfd, filename, name_len, addr, &addr_len) < 0) 
    {
        perror("recv filename error");
        goto cleanup;
    }
    
    printf("\n%s receiving.....\n", filename);
    
    // 接收文件大小
    uint64_t net_file_size;
    if (recv_all(sockfd, &net_file_size, sizeof(net_file_size), addr, &addr_len) < 0) 
    {
        perror("recv file size error");
        goto cleanup;
    }
    
    uint64_t file_size = be64toh(net_file_size); // 转换字节序   
    printf("File size: %lu bytes\n", file_size);
    
    if (file_size > 10 * 1024 * 1024 * 1024ULL) // 检查文件大小是否合理 (10GB)
    { 
        fprintf(stderr, "File size too large: %lu bytes\n", file_size);
        goto cleanup;
    }
    
    // 构建保存路径
    char filepath[512] = {0};
    // 确保保存目录不以斜杠结尾
    char clean_save_dir[256];
    strncpy(clean_save_dir, save_dir, sizeof(clean_save_dir)-1);
    clean_save_dir[sizeof(clean_save_dir)-1] = '\0';
    
    // 移除末尾的斜杠
    size_t len = strlen(clean_save_dir);
    if (len > 0 && clean_save_dir[len-1] == '/') 
    {
        clean_save_dir[len-1] = '\0';
    }
    
    snprintf(filepath, sizeof(filepath), "%s/%s", clean_save_dir, filename);
    printf("Saving to: %s\n", filepath);
    
    // 确保目录存在
    char *dir_path = strdup(filepath);
    char *dir_name = dirname(dir_path);
    struct stat st = {0};
    if (stat(dir_name, &st) == -1) 
    {
        mkdir(dir_name, 0700);
    }
    free(dir_path);
    
    // 打开文件
    file = fopen(filepath, "wb");
    if (!file) 
    {
        fprintf(stderr, "fopen error for %s: %s\n", filepath, strerror(errno));
        goto cleanup;
    }
    
    // 接收文件内容
    char recvbuff[1400]; // UDP适合较小的数据包
    uint64_t total_received = 0;
    
    while (total_received < file_size) 
    {
        size_t remaining = file_size - total_received;
        size_t to_receive = remaining > sizeof(recvbuff) ? sizeof(recvbuff) : remaining;
        
        ssize_t recvlen = recvfrom(sockfd, recvbuff, to_receive, 0, (struct sockaddr*)addr, &addr_len);
        if (recvlen <= 0) 
        {
            perror("recv file content error");
            goto cleanup;
        }
        
        total_received += recvlen;
        if (fwrite(recvbuff, 1, recvlen, file) != recvlen) 
        {
            perror("fwrite error");
            goto cleanup;
        }
        
        // 显示进度
        printf("Received: %lu/%lu bytes (%.2f%%)\r", 
               total_received, file_size, (double)total_received/file_size*100);
        fflush(stdout);
    }
    
    printf("\n%s received successfully!\n", filename);
    ret = 0; 

cleanup:
    if (file) 
    {
        fclose(file);
    }
    return ret;
}

// 线程函数
void* file_transfer_thread(void *arg) 
{
    thread_arg_t *thread_arg = (thread_arg_t*)arg;
    int result;

    if (thread_arg->mIs_send) 
    {
        result = send_file(thread_arg->msockfd, thread_arg->mfilepath, 
                          &thread_arg->maddr, thread_arg->maddr_len);
    } 
    else 
    {
        result = recv_file(thread_arg->maddr_len, thread_arg->msave_dir, 
                          &thread_arg->maddr, thread_arg->maddr_len);
    }
    
    if (result != 0) 
    {
        fprintf(stderr, "File transfer failed in thread\n");
    }

    return NULL;
}